# --------------------------------------------#
#   该部分代码用于看网络结构
# --------------------------------------------#
import torch
import configparser
from thop import clever_format, profile
from torchsummary import summary
from utils.utils import get_classes
from nets.yolo import YoloBody

# 解析参数
content = configparser.ConfigParser()
content.read("config.ini", encoding='UTF-8')

classes_path = content.get("train", "classes_path")

attention_type = eval(content.get("train", "attention_type"))
enable_spp = eval(content.get("train", "enable_spp"))
enable_mobile = eval(content.get("train", "enable_mobile"))


def view_model():
    class_names, num_classes = get_classes(classes_path)

    input_shape = [416, 416]
    anchors_mask = [[3, 4, 5], [1, 2, 3]]
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    m = YoloBody(anchors_mask, num_classes, attention_type, enable_spp, enable_mobile).to(device)

    summary(m, (3, input_shape[0], input_shape[1]))

    dummy_input = torch.randn(1, 3, input_shape[0], input_shape[1]).to(device)
    flops, params = profile(m.to(device), (dummy_input,), verbose=False)
    # --------------------------------------------------------#
    #   flops * 2是因为profile没有将卷积作为两个operations
    #   有些论文将卷积算乘法、加法两个operations。此时乘2
    #   有些论文只考虑乘法的运算次数，忽略加法。此时不乘2
    #   本代码选择乘2，参考YOLOX。
    # --------------------------------------------------------#
    flops = flops * 2
    flops, params = clever_format([flops, params], "%.3f")
    # noinspection PyRedundantParentheses
    print('Total GFLOPS: %s' % (flops))
    # noinspection PyRedundantParentheses
    print('Total params: %s\n' % (params))


if __name__ == "__main__":
    view_model()
